# standard module
import sys
from enum import Enum
from pathlib import Path
from datetime import datetime
from contextlib import asynccontextmanager

from tqdm import tqdm
from sqlalchemy import create_engine
import pandas as pd

# Data Source
import akshare as ak


from models.model import StockData
from db_utils import engine

class DataOp(Enum):
    stock_info_to_local_file = "get_stock_info"
    stock_data_to_local_file = "get_hist_data"
    hist_data_to_database = "store_hist_data"


DATA_FIEL_DIR = "data_files"
DF_TO_DATABASE_MAPPING = {
    "日期": "date",
    "股票代码": "stock_code",
    "开盘": "open",
    "收盘": "close",
    "最高": "high",
    "最低": "low",
    "成交量": "volume",
    "成交额": "turnover_amount",
    "振幅": "amplitude",
    "涨跌幅": "change_percent",
    "涨跌额": "change_amount",
    "换手率": "turnover_rate"
}

def stock_data_to_local_file(stock_symbols:list[str],
                             start_date:str = "20100101",
                             end_date:str = "20250620",
                             period:str = "daily",
                             adjust:str = "qfq",
                             output_dir:str = DATA_FIEL_DIR):
    # TODO: figure out datetime format stuff
    try:
        date_format= "%Y%m%d"
        _start = None
        _end= None
        try:
            _start = datetime.strptime(start_date,date_format)
            _end = datetime.strptime(end_date,date_format)
        except:
            raise ValueError("Invalid date format")

        if _start > _end:
            raise ValueError("Invalid date value : start_date can't be greater the end_date!")
    except ValueError as e:
        print(e)

    try:
        Path(output_dir).mkdir(parents=True,exist_ok=True)
        for symbol in stock_symbols:
            file_name = f"{symbol}_{adjust}_{start_date}_{end_date}.csv"
            file_path = Path(output_dir)/file_name
            if file_path.exists():
                print(f"{file_path} exists!")
                continue

            df = ak.stock_zh_a_hist(
                symbol=symbol,
                period=period,
                start_date=start_date,
                end_date= end_date,
                adjust="qfq"
            )

            if not df.empty:
                df.to_csv(Path(output_dir)/file_name)
                print(f"Success fetching {symbol}!")
            else:
                print(f"Error when fetching {symbol}!")
    except:
        raise RuntimeError("Error when fetching stock history data!")

def stock_data_to_database(df:pd.DataFrame):
    if not df.empty:
        df_renamed = df.rename(columns=DF_TO_DATABASE_MAPPING)
        df_renamed.to_sql(
            name=StockData.__tablename__,
            con = engine,
            if_exists="append",
            index=False,
            chunksize=3000,
            method="multi",            
        )
    
    pass

def read_file_hist_data(stock_code:str):
    file_name_gen = (item.name for item in Path(DATA_FIEL_DIR).iterdir() 
                    if item.name.startswith(stock_code))
    file_name = next(file_name_gen,None)
    hist_data_df = pd.DataFrame()
    if file_name:
        hist_data_df = pd.read_csv(Path(DATA_FIEL_DIR)/file_name,index_col=0,dtype={"股票代码":str})        
        print(file_name + ":")
        print(hist_data_df.head())
    else:
        print(f"Can't find stock data of {stock_code}")
    return hist_data_df

def stock_info_to_local_file():
    try:
        file_path = Path(DATA_FIEL_DIR)/"all_stock_code_a.csv"
        all_stock_code_a_df = ak.stock_info_a_code_name()
        all_stock_code_a_df.to_csv(file_path)
        print(f"All Stock Code Saved To {file_path}")
    except:
        print("Error When saving All Stock Code")


def cli():
    if len(sys.argv) < 2:
        print(f"Usage: {sys.argv[0]} [operation] [stock_symbols]"
                "(start_date) (end_date)"
                "(period) (adjust) (output_dir)")
        sys.exit(1)

    print(sys.argv)
    if sys.argv[1] == DataOp.stock_info_to_local_file.value:
        stock_info_to_local_file()
    elif sys.argv[1] == DataOp.stock_data_to_local_file.value:
        symbols =[]
        symbols = sys.argv[2].split(",")
        # print(sys.argv)
        if len(sys.argv) == 3:
            stock_data_to_local_file(stock_symbols=symbols)
        elif len(sys.argv) == 5:
            stock_data_to_local_file(stock_symbols=symbols,
                                    start_date=sys.argv[2],
                                    end_date=sys.argv[3])
    elif sys.argv[1] == DataOp.hist_data_to_database.value:
        df = read_file_hist_data(sys.argv[2])
        stock_data_to_database(df)


if __name__ == '__main__':
    cli()